import cv2
import torch

from cjc import project_setting
from cjc.detector.detecotor_single import DetectorSingle
from cjc.detector.opt_tool import create_opt


class SexDetector(DetectorSingle):
    def __init__(self, _opt):
        _opt.weights = project_setting.root+'/cjc/weights/smile.pt'
        _opt.source = '-1'
        _opt.name = 'smile'
        # _opt.view_img = True
        super(SexDetector, self).__init__(_opt)



if __name__ == '__main__':
    opt = create_opt()
    opt.view_img =True
    detector = SexDetector(opt)
    cv_cap = cv2.VideoCapture(0)
    while True:
        ret, frame = cv_cap.read()
        if ret:
            with torch.no_grad():
                a = detector.detect(frame,0)
                cv2.waitKey(1)
                print(a)
